{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "3tW2MkdrIn6F",
      "metadata": {
        "id": "3tW2MkdrIn6F"
      },
      "source": [

        "\n",
        "## Installation\n",
        "\n",
        "To run Jupyter notebook locally:\n",
        "\n",
        "```\n",
        "python3 -m pip install \"waymo_open_dataset_tf_2_6_0==1.4.3\"\n",
        "python3 -m pip install \"notebook\u003e=5.3\" \"ipywidgets\u003e=7.5\"\n",
        "python3 -m pip install --upgrade jupyter_http_over_ws\u003e=0.0.7 \u0026\u0026 \\\n",
        "jupyter serverextension enable --py jupyter_http_over_ws\n",
        "jupyter notebook\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "xdEsPZX9LCVo",
      "metadata": {
        "id": "xdEsPZX9LCVo"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "As_l5B3o8sl6",
      "metadata": {
        "cellView": "code",
        "id": "As_l5B3o8sl6"
      },
      "outputs": [],
      "source": [
        "#@title Waymo Open Dataset imports\n",
        "from waymo_open_dataset import dataset_pb2\n",
        "from waymo_open_dataset import label_pb2\n",
        "from waymo_open_dataset.metrics.python import keypoint_metrics\n",
        "from waymo_open_dataset.protos import keypoint_pb2\n",
        "from waymo_open_dataset.utils import box_utils\n",
        "from waymo_open_dataset.utils import frame_utils\n",
        "from waymo_open_dataset.utils import keypoint_data\n",
        "from waymo_open_dataset.utils import keypoint_draw\n",
        "from waymo_open_dataset.utils import range_image_utils\n",
        "from waymo_open_dataset.utils import transform_utils"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c_kFiqySorjE",
      "metadata": {
        "id": "c_kFiqySorjE"
      },
      "outputs": [],
      "source": [
        "# File path to a tfrecods file with Frame protos with human keypoints.\n",
        "frame_path = 'frame_with_keypoints.tfrecord'\n",
        "\n",

        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6f2d89b6",
      "metadata": {
        "cellView": "code",
        "id": "6f2d89b6"
      },
      "outputs": [],
      "source": [
        "#@title Load Frame proto\n",
        "import tensorflow as tf\n",
        "\n",
        "dataset = tf.data.TFRecordDataset(frame_path, compression_type='')\n",
        "for data in dataset:\n",
        "  frame = dataset_pb2.Frame()\n",
        "  frame.ParseFromString(bytearray(data.numpy()))\n",
        "  break\n",
        "\n",
        "labels = keypoint_data.group_object_labels(frame)\n",
        "print(f'Loaded {len(labels)} objects')"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "_FJKloajRyMt",
      "metadata": {
        "id": "_FJKloajRyMt"
      },
      "source": [
        "## Visualization"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6mlBUrd9-g3k",
      "metadata": {
        "id": "6mlBUrd9-g3k"
      },
      "outputs": [],
      "source": [
        "#@title Auixiliary imports and utils\n",
        "\n",
        "import os\n",
        "import math\n",
        "import numpy as np\n",
        "from matplotlib import pylab as plt\n",
        "import plotly.graph_objects as go\n",
        "import itertools\n",
        "import PIL.Image\n",
        "import io\n",
        "import dataclasses\n",
        "\n",
        "\n",
        "def _imdecode(buf: bytes) -\u003e np.ndarray:\n",
        "  with io.BytesIO(buf) as fd:\n",
        "    pil = PIL.Image.open(fd)\n",
        "    return np.array(pil)\n",
        "\n",
        "\n",
        "def _imshow(ax: plt.Axes, image_np: np.ndarray):\n",
        "  ax.imshow(image_np)\n",
        "  ax.axis('off')\n",
        "  ax.set_autoscale_on(False)\n",
        "\n",
        "\n",
        "def _draw_laser_points(fig: go.Figure,\n",
        "                       points: np.ndarray,\n",
        "                       color: str = 'gray',\n",
        "                       size: int = 3):\n",
        "  \"\"\"Visualizes laser points on a plotly figure.\"\"\"\n",
        "  fig.add_trace(\n",
        "      go.Scatter3d(\n",
        "          mode='markers',\n",
        "          x=points[:, 0],\n",
        "          y=points[:, 1],\n",
        "          z=points[:, 2],\n",
        "          marker=dict(color=color, size=size)))\n",
        "\n",
        "\n",
        "def _create_plotly_figure() -\u003e go.Figure:\n",
        "  \"\"\"Creates a plotly figure for 3D visualization.\"\"\"\n",
        "  fig = go.Figure()\n",
        "  axis_settings = dict(\n",
        "      showgrid=False,\n",
        "      zeroline=False,\n",
        "      showline=False,\n",
        "      showbackground=False,\n",
        "      showaxeslabels=False,\n",
        "      showticklabels=False)\n",
        "  fig.update_layout(\n",
        "      width=600,\n",
        "      height=600,\n",
        "      showlegend=False,\n",
        "      scene=dict(\n",
        "          aspectmode='data',  # force xyz has same scale,\n",
        "          xaxis=axis_settings,\n",
        "          yaxis=axis_settings,\n",
        "          zaxis=axis_settings,\n",
        "      ),\n",
        "  )\n",
        "  return fig"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "RiBlI9LK2fKs",
      "metadata": {
        "id": "RiBlI9LK2fKs"
      },
      "outputs": [],
      "source": [
        "#@title Select object and camera\n",
        "object_id = 'DQFLdFau_A8kTPOkDxfgJA'\n",
        "camera_name = dataset_pb2.CameraName.Name.FRONT_RIGHT\n",
        "\n",
        "camera_image_by_name = {i.name: i.image for i in frame.images}\n",
        "obj = labels[object_id]\n",
        "num_laser_points = len(obj.laser.keypoints.keypoint)\n",
        "num_camera_points = len(obj.camera[camera_name].keypoints.keypoint)\n",
        "\n",
        "print(f'Object {object_id} has')\n",
        "print(f'{num_laser_points} laser keypoints '\n",
        "      '(short name | location | is_occluded):')\n",
        "for k in sorted(obj.laser.keypoints.keypoint, key=lambda k: k.type):\n",
        "  m = k.keypoint_3d.location_m\n",
        "  location_str = f'({m.x:.2f}, {m.y:.2f}, {m.z:.2f})'\n",
        "  print(f'{keypoint_draw.point_name(k.type)}\\t|'\n",
        "        f' {location_str:25} | {k.keypoint_3d.visibility.is_occluded}')\n",
        "print(f'\\na LaserKeypoint proto example:\\n\\n{obj.laser.keypoints.keypoint[0]}')\n",
        "\n",
        "print(f'{num_camera_points} camera keypoints '\n",
        "      '(short name |  location | is_occluded):')\n",
        "for k in sorted(\n",
        "    obj.camera[camera_name].keypoints.keypoint, key=lambda k: k.type):\n",
        "  px = k.keypoint_2d.location_px\n",
        "  location_str = f'({px.x:.0f}, {px.y:.0f})'\n",
        "  print(f'{keypoint_draw.point_name(k.type)}\\t'\n",
        "        f'| {location_str:13} | {k.keypoint_2d.visibility.is_occluded}')\n",
        "print(f'\\na CameraKeypoint proto example:\\n\\n'\n",
        "      f'{obj.camera[camera_name].keypoints.keypoint[0]}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "Zya-PvdW2fKt",
      "metadata": {
        "id": "Zya-PvdW2fKt"
      },
      "outputs": [],
      "source": [
        "#@title Show camera keypoints\n",
        "image_np = _imdecode(camera_image_by_name[camera_name])\n",
        "croped_image, cropped_camera_keypoints = keypoint_draw.crop_camera_keypoints(\n",
        "    image_np,\n",
        "    obj.camera[camera_name].keypoints.keypoint,\n",
        "    obj.camera[camera_name].box,\n",
        "    margin=0.3)\n",
        "camera_wireframe = keypoint_draw.build_camera_wireframe(\n",
        "    cropped_camera_keypoints)\n",
        "\n",
        "keypoint_draw.OCCLUDED_BORDER_WIDTH = 3\n",
        "_, ax = plt.subplots(frameon=False, figsize=(5, 7))\n",
        "_imshow(ax, croped_image)\n",
        "keypoint_draw.draw_camera_wireframe(ax, camera_wireframe)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "Yc-erhA72fKy",
      "metadata": {
        "id": "Yc-erhA72fKy"
      },
      "outputs": [],
      "source": [
        "#@title Show laser keypoints\n",
        "\n",
        "# Select laser points inside pedestrian's bounding box\n",
        "(range_images, camera_projections, _, range_image_top_pose\n",
        ") = frame_utils.parse_range_image_and_camera_projection(frame)\n",
        "points, cp_points = frame_utils.convert_range_image_to_point_cloud(\n",
        "    frame, range_images, camera_projections, range_image_top_pose)\n",
        "points_all = np.concatenate(points, axis=0)\n",
        "box = box_utils.box_to_tensor(obj.laser.box)[tf.newaxis, :]\n",
        "box_points = points_all[box_utils.is_within_box_3d(points_all, box)[:, 0]]\n",
        "print(f'{box_points.shape[0]} laser points selected.')\n",
        "\n",
        "# Visualize 3D scene\n",
        "laser_wireframe = keypoint_draw.build_laser_wireframe(\n",
        "    obj.laser.keypoints.keypoint)\n",
        "fig = _create_plotly_figure()\n",
        "keypoint_draw.draw_laser_wireframe(fig, laser_wireframe)\n",
        "_draw_laser_points(fig, box_points)\n",
        "fig.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "bz4MaFMTRnh_",
      "metadata": {
        "id": "bz4MaFMTRnh_"
      },
      "source": [
        "## Metrics\n",
        "\n",
        "`waymo_open_dataset` package supports the following metrics:\n",
        "\n",
        "- Precision (P) and average precision (AP) for various thresholds of Object Keypoint Similarity (OKS).\n",
        "- Percentage of Correct Keypoints (PCK)\n",
        "- Mean Per Joint Position Error MPJPE (aka MPJE)\n",
        "\n",
        "Please refer to [`waymo_open_dataset/metrics/python/keypoint_metrics.py`](https://github.com/waymo-research/waymo-open-dataset/blob/master/waymo_open_dataset/metrics/keypoint_metrics.py) for implementation details.\n",
        "\n",
        "NOTE: The OKS metric penalizes keypoint coordinates far outside the ground truth bounding box for samples without ground truth keypoints. For such cases we set the default coordinates to be in the middle of the object's box to avoid the penalty."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "W3rreupyhnL_",
      "metadata": {
        "id": "W3rreupyhnL_"
      },
      "outputs": [],
      "source": [
        "#@title Example how to compute metrics for camera keypoints\n",
        "from typing import Tuple\n",
        "\n",
        "def get_camera_data(\n",
        "    frame: dataset_pb2.Frame\n",
        ") -\u003e Tuple[keypoint_data.KeypointsTensors, keypoint_data.KeypointsTensors]:\n",
        "  \"\"\"Extracts camera keypoints and bounding boxes from the input Frame proto.\"\"\"\n",
        "  all_keypoints = []\n",
        "  all_boxes = []\n",
        "  for cl in frame.camera_labels:\n",
        "    for l in cl.labels:\n",
        "      if l.HasField('camera_keypoints'):\n",
        "        box = keypoint_data.create_camera_box_tensors(l.box, dtype=tf.float32)\n",
        "        keypoints = keypoint_data.create_camera_keypoints_tensors(\n",
        "            l.camera_keypoints.keypoint,\n",
        "            default_location=box.center,\n",
        "            order=keypoint_data.CANONICAL_ORDER_CAMERA,\n",
        "            dtype=tf.float32)\n",
        "        all_keypoints.append(keypoints)\n",
        "        all_boxes.append(box)\n",
        "  keypoint_tensors = keypoint_data.stack_keypoints(all_keypoints)\n",
        "  box_tensors = keypoint_data.stack_boxes(all_boxes)\n",
        "  return keypoint_tensors, box_tensors\n",
        "\n",
        "\n",
        "gt_cam, gt_cam_box = get_camera_data(frame)\n",
        "\n",
        "noise_stddev = 5.0  # in pixels\n",
        "pr_cam = keypoint_data.KeypointsTensors(\n",
        "    location=gt_cam.location +\n",
        "    tf.random.normal(gt_cam.location.shape, stddev=noise_stddev),\n",
        "    visibility=gt_cam.visibility)\n",
        "\n",
        "all_metrics = keypoint_metrics.create_combined_metric(\n",
        "    keypoint_metrics.DEFAULT_CONFIG_CAMERA)\n",
        "all_metrics.update_state([gt_cam, pr_cam, gt_cam_box])\n",
        "result = all_metrics.result()\n",
        "\n",
        "print('Camera keypoint metrics:')\n",
        "for name, tensor in sorted(result.items(), key=lambda e: e[0]):\n",
        "  print(f'{name:20s}: {tensor.numpy():.3f}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ZibXsMaXPRrL",
      "metadata": {
        "id": "ZibXsMaXPRrL"
      },
      "outputs": [],
      "source": [
        "#@title Example how to compute metrics for laser keypoints\n",
        "\n",
        "\n",
        "def get_laser_data(\n",
        "    frame: dataset_pb2.Frame\n",
        ") -\u003e Tuple[keypoint_data.KeypointsTensors, keypoint_data.KeypointsTensors]:\n",
        "  \"\"\"Extracts laser keypoints and bounding boxes from the input Frame proto.\"\"\"\n",
        "  all_keypoints = []\n",
        "  all_boxes = []\n",
        "  for l in frame.laser_labels:\n",
        "    if l.HasField('laser_keypoints'):\n",
        "      box = keypoint_data.create_laser_box_tensors(l.box, dtype=tf.float32)\n",
        "      keypoints = keypoint_data.create_laser_keypoints_tensors(\n",
        "          l.laser_keypoints.keypoint,\n",
        "          default_location=box.center,\n",
        "          order=keypoint_data.CANONICAL_ORDER_LASER,\n",
        "          dtype=tf.float32)\n",
        "      all_keypoints.append(keypoints)\n",
        "      all_boxes.append(box)\n",
        "  keypoint_tensors = keypoint_data.stack_keypoints(all_keypoints)\n",
        "  box_tensors = keypoint_data.stack_boxes(all_boxes)\n",
        "  return keypoint_tensors, box_tensors\n",
        "\n",
        "\n",
        "gt_cam, gt_cam_box = get_laser_data(frame)\n",
        "\n",
        "noise_stddev = 0.05  # in meters\n",
        "pr_cam = keypoint_data.KeypointsTensors(\n",
        "    location=gt_cam.location +\n",
        "    tf.random.normal(gt_cam.location.shape, stddev=noise_stddev),\n",
        "    visibility=gt_cam.visibility)\n",
        "\n",
        "all_metrics = keypoint_metrics.create_combined_metric(\n",
        "    keypoint_metrics.DEFAULT_CONFIG_LASER)\n",
        "all_metrics.update_state([gt_cam, pr_cam, gt_cam_box])\n",
        "result = all_metrics.result()\n",
        "\n",
        "print('Laser keypoint metrics:')\n",
        "for name, tensor in sorted(result.items(), key=lambda e: e[0]):\n",
        "  print(f'{name:20s}: {tensor.numpy():.3f}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3vxkcm3GqQ2w",
      "metadata": {
        "id": "3vxkcm3GqQ2w"
      },
      "outputs": [],
      "source": [
        "#@title Use individual metrics\n",
        "\n",
        "per_type_scales = [\n",
        "    keypoint_metrics.DEFAULT_PER_TYPE_SCALES[t]\n",
        "    for t in keypoint_data.CANONICAL_ORDER_CAMERA\n",
        "]\n",
        "oks = keypoint_metrics.AveragePrecisionAtOKS(per_type_scales, thresholds=[0.95])\n",
        "oks.update_state([gt_cam, pr_cam, gt_cam_box])\n",
        "oks.result()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "",
        "kind": "local"
      },
      "name": "Waymo Open Dataset (Perception): Human Keypoints Tutorial",
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "tutorial_keypoints.ipynb",
          "timestamp": 1644014756220
        },
        {
          "file_id": "tutorial_keypoints.ipynb",
          "timestamp": 1643966716329
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.9.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}
